use anyhow::Result;
use anyhow::anyhow;
use bytes::Buf;
use bytes::BytesMut;
use std::net::Ipv4Addr;
use std::net::Ipv6Addr;
use std::net::SocketAddr;
use std::net::SocketAddrV4;
use std::net::SocketAddrV6;
use std::str::FromStr;
use std::sync::Arc;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};

use crate::proxy::InOutBound;
use crate::proxy::Proxy;
use crate::proxy_config::InBoundKind;
use crate::proxy_config::ProxyConfig;

use super::socks_proto::SOCKS_VERSION;
use super::socks_proto::SUPPORT_AUTHS;
use super::socks_proto::SocksAuthMethod;
use super::socks_proto::SocksRequestCommand;
use super::socks_proto::SocksResponse;
use super::socks_proto::SocksResponseAddressType;
use super::socks_proto::SocksResponseType;

const BUFF_SIZE: usize = 10240;

pub struct SocksServer {
    tcp_listener: Option<TcpListener>,

    flag: InOutBound,
    config: Arc<ProxyConfig>,
}


impl Proxy for SocksServer {
    async fn start_serve(&mut self) -> Result<()> {
        let listen_addr =
            SocketAddr::from_str("127.0.0.1:1080").expect("parse to socket addr failed!");
        let tcplistener = TcpListener::bind(listen_addr).await?;
        log::info!("listen 1080 success!");
        self.tcp_listener = Some(tcplistener);

        loop {
            if let Ok((tcpstream, addr)) = self.tcp_listener.as_ref().unwrap().accept().await {
                tokio::spawn(async move {
                    let server = Socks5SessionHandler::new(tcpstream);
                    let _ = server.handle_request(addr).await;
                });
            }
        }
    }

    async fn prepare_serve(
        config: std::sync::Arc<crate::proxy_config::ProxyConfig>,
        flag: InOutBound,
    ) -> Result<Self>
    where
        Self: Sized,
    {
        if flag == InOutBound::OutBound {
            return Err(anyhow!("Socks不能做为OutBound使用"));
        }
        if config.in_bound.kind != InBoundKind::Socks5 {
            return Err(anyhow!("InBound kind 必须是socks5"));
        }

        let s = Self {
            tcp_listener: None,
            config,
            flag,
        };
        Ok(s)
    }
}

fn find_both_support_auth_method(
    client_auth_methods: &[SocksAuthMethod],
    server_auth_methods: &[SocksAuthMethod],
) -> Vec<SocksAuthMethod> {
    server_auth_methods
        .iter()
        .filter(|sm| client_auth_methods.contains(sm))
        .copied()
        .collect()
}

fn read_address(bs: &mut BytesMut) -> Result<(SocksResponseAddressType, Vec<u8>, u16)> {
    if bs.len() < 3 {
        return Err(anyhow!("SocksRequestAddressType的有效长度不可能小于2"));
    }
    let t = &bs.get_u8();
    match t {
        0x01 => Ok((
            SocksResponseAddressType::Ipv4,
            bs.split_to(4).to_vec(), //todo: check to le_bytes是否正确
            bs.get_u16(),
        )),
        0x02 => {
            let len = bs.get_u8() as usize;
            let domain = bs.split_to(len);
            Ok((
                SocksResponseAddressType::Domain,
                domain.to_vec(),
                bs.get_u16(),
            ))
        }
        0x03 => Ok((
            SocksResponseAddressType::Ipv6,
            bs.split_to(16).to_vec(),
            bs.get_u16(),
        )),
        _ => Err(anyhow!("unkonw atyp")),
    }
}

struct Socks5SessionHandler {
    request_stream: TcpStream,
    client_support_authmethods: Option<Vec<SocksAuthMethod>>,
    selected_auth_method: Option<SocksAuthMethod>,
    response: Option<SocksResponse>,
    outbound_stream: Option<TcpStream>,
    buffer: BytesMut,
}
impl Socks5SessionHandler {
    fn new(request: TcpStream) -> Self {
        Self {
            request_stream: request,
            client_support_authmethods: None,
            selected_auth_method: None,
            response: None,
            outbound_stream: None,
            buffer: BytesMut::with_capacity(1024 * 5),
        }
    }
    async fn handle_request(mut self, _src_addr: SocketAddr) -> Result<()> {
        self.auth().await?;
        self.method_request().await?;
        self.data_transfer().await?;
        Ok(())
    }
    pub async fn auth(&mut self) -> Result<()> {
        let n = self.request_stream.read_buf(&mut self.buffer).await?;
        log::trace!("auth nego read {} size request", n);
        let ver = self.buffer.get_u8();
        log::info!("socks ver is: {}", ver);
        if ver != SOCKS_VERSION {
            return Err(anyhow!("socks version is incorrect"));
        }
        let methods_len = self.buffer.get_u8();
        log::info!("client methods lens is:{}", methods_len);
        let auth_methods = self.buffer.split_to(methods_len as usize);

        log::info!("{:?}", auth_methods);
        let client_auth_methods: Vec<_> = auth_methods
            .iter()
            .map(|x| SocksAuthMethod::from(*x))
            .collect();
        self.client_support_authmethods = Some(client_auth_methods);
        let mut both_support_am = find_both_support_auth_method(
            self.client_support_authmethods.as_ref().unwrap(),
            &SUPPORT_AUTHS,
        );
        both_support_am.sort();
        let auth_method = both_support_am
            .first()
            .expect("there is must be an authmethod");
        self.selected_auth_method = Some(*auth_method);
        //写回选定的认证方式
        self.request_stream
            .write_all(&[SOCKS_VERSION, (*auth_method).into()])
            .await?;

        //将buffsize重置
        self.buffer.reserve(BUFF_SIZE);
        Ok(())
    }

    pub async fn method_request(&mut self) -> Result<()> {
        let n = self.request_stream.read_buf(&mut self.buffer).await?;
        log::trace!("had read {} to in method_request", n);
        log::trace!("starting parse to socks request");
        let ver = self.buffer.get_u8();
        if ver != SOCKS_VERSION {
            return Err(anyhow!("socks的版本不对"));
        }
        let cmd = self.buffer.get_u8();
        let mut response_builder = SocksResponse::builder();
        match SocksRequestCommand::try_from(cmd) {
            Ok(cmd) => {
                let _rsv = self.buffer.get_u8();
                let (atyp, address, port) = read_address(&mut self.buffer)?;
                response_builder.atyp(atyp);

                response_builder.dst_port(port);
                let ipaddr: Option<SocketAddr> = match atyp {
                    SocksResponseAddressType::Ipv4 => {
                        let ip = Ipv4Addr::new(address[0], address[1], address[2], address[3]);
                        response_builder.dst_addr(address);
                        let addr = SocketAddrV4::new(ip, port);
                        Some(addr.into())
                    }

                    SocksResponseAddressType::Domain => {
                        //解析domain为ip地址
                        todo!()
                    }
                    SocksResponseAddressType::Ipv6 => {
                        let t: Result<[u8; 16]> = address
                            .try_into()
                            .map_err(|_e| anyhow!("address convert to [u8:16] faild!"));
                        match t {
                            Ok(t) => {
                                let ip = Ipv6Addr::from(t);
                                let addr = SocketAddrV6::new(ip, port, 0, 0);
                                Some(addr.into())
                            }
                            Err(_e) => {
                                response_builder.rep(SocksResponseType::AddressIncorrect);
                                None
                            }
                        }
                    }
                };
                if let Some(socket_addr) = ipaddr {
                    match cmd {
                        SocksRequestCommand::Connect => match TcpStream::connect(socket_addr).await
                        {
                            Ok(stream) => {
                                response_builder.rep(SocksResponseType::Success);
                                self.outbound_stream = Some(stream);
                            }
                            Err(e) => {
                                log::error!("connect failed! {}", e);
                                response_builder.rep(SocksResponseType::NetWorkInAccessible);
                            }
                        },
                        SocksRequestCommand::Bind => todo!(),
                        SocksRequestCommand::Udp => todo!(),
                    }
                }
            }
            Err(_e) => {
                log::error!("not support command!");
                response_builder.rep(SocksResponseType::NoSupportCommand);
            }
        }

        self.buffer.reserve(BUFF_SIZE);
        self.response = Some(response_builder.build());
        self.request_stream
            .write_all(&self.response.as_ref().unwrap().to_bytes())
            .await?;
        Ok(())
    }

    pub async fn data_transfer(&mut self) -> Result<()> {
        if let Some(mut server_stream) = self.outbound_stream.as_mut()
            && self.response.as_ref().unwrap().rep() == SocksResponseType::Success
        {
            let (s1, s2) =
                tokio::io::copy_bidirectional(&mut self.request_stream, &mut server_stream).await?;
            log::info!("对考结束，s1:{},s2:{}", s1, s2);
            Ok(())
        } else {
            Err(anyhow!("handle request failed!"))
        }
    }
}

#[cfg(test)]
mod test {
    use super::super::socks_proto::SocksResponseAddressType;
    use super::*;
    use bytes::BufMut;

    #[test]
    fn read_address_test() {
        let mut bs = BytesMut::new();
        bs.put_u8(0x01);
        bs.put_u8(192);
        bs.put(&[168, 55, 100][..]);
        bs.put_u16(8000);

        match read_address(&mut bs) {
            Err(e) => {
                panic!("{}", e);
            }
            Ok((aty, address, port)) => {
                assert_eq!(aty, SocksResponseAddressType::Ipv4);
                assert_eq!(address, vec![192, 168, 55, 100]);
                assert_eq!(port, 8000);
            }
        }
    }
}
